from typing import TypedDict
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph
from src.common.logger import getLogger
from langgraph.constants import START, END
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class StockAct(BaseModel):
    symbol: str = Field(description = "沪深300指数的股票代码")
    date: str = Field(description = "查询的时间")

class StockState(TypedDict):
    query: str
    symbol: str
    date: str
    answer: str

class StockAgent:

    def __init__(self, llm_model, tool):
        self.llm_model = llm_model
        self.tool = tool

    def extract_node(self, state: StockState):
        logger.info("StockAgent extract_node start")
        query = state["query"]
        extract_template = """
            你是一位问题分析高手，根据用户的输入问题提取其中的股票代码和查询时间。
            用户输入的问题：{question}
        """
        extract_prompt = ChatPromptTemplate.from_template(extract_template)
        extract_chain = extract_prompt | self.llm_model.with_structured_output(StockAct)
        extract_result = extract_chain.invoke({ "question": query })
        logger.info(f"StockAgent extract_node extract_result: {extract_result}")
        return { "symbol": extract_result.symbol, "date": extract_template.date }

    def execute_node(self, state: StockState):
        logger.info("StockAgent execute_node start")
        param = { "symbol": state["symbol"], "date": state["date"] }
        result = self.tool.func(*param)
        logger.info(f"StockAgent execute_node result len: {len(result)}")
        return { "answer": result }

    def build_graph(self):
        logger.info("StockAgent build_graph start")
        graph = StateGraph(StockAgent)
        graph.add_node("extract", self.extract_node)
        graph.add_node("execute", self.execute_node)

        graph.add_edge(START, "extract")
        graph.add_edge("extract", "execute")
        graph.add_edge("execute", END)
        return graph.compile()

    def invoke(self, query):
        logger.info(f"StockAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        logger.info(f"StockAgent invoke response: {response}")
        return { "document": response.get("answer", None), "answer": response.get("answer", None) }
